import torch
import torch.nn as nn

class UnitGaussianNormalizer(nn.Module):
    def __init__(self, mean, std, eps=1e-05):
        super().__init__()
        self.mean = torch.tensor(mean)
        self.std = torch.tensor(std)
        self.eps = torch.tensor(eps)

    def encode(self, x):
        x = (x - self.mean) / (self.std + self.eps)
        return x

    def decode(self, x):
        x = x * (self.std.to(x) + self.eps.to(x)) + self.mean.to(x)
        return x

    def to(self,device):
        self.mean = self.mean.to(device)
        self.std  = self.std .to(device)
        self.eps  = self.eps .to(device)


